import sys
import os

current_dir = os.path.dirname(os.path.abspath(__file__))
grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(grandparent_dir)

import random
import numpy as np
import time
import pickle
from object.Buiding import Building
import matplotlib.pyplot as plt
random.seed(int(time.time()))
# 参数区
num_floors = 5          # 楼层数
floor_width = 100       # 每层宽（格子数）
floor_length = 100      # 每层长（格子数）
floor_height = 5      # 楼梯间高度 (m)
stairs_per_floor = 5 # 每层上楼梯数量
exits_on_ground = 5     # 一楼出口数量
cell_size = 1.0         # 每格面积(m²)
obstacle_size = (3,15)   # 障碍物尺寸
obstacle_attempt = 10    # 放障碍物尝试 (2*4) 或 (4*2) 形状
speed_floor_down_radio = 1.0 #下楼速度和平路运动的比值
time_update = 1     # 每帧1秒 (每1秒更新一次)
wall_thickness = 3  # 墙厚度
staircase_width = 2 # 楼梯口宽度
exit_width = 2      # 出口宽度

# 建筑位置状态
# 0表示空地
# 1表示障碍物或墙
# 2表示上楼梯位置
# 3表示下楼梯位置
# 4表示出口

floors = np.zeros((num_floors, floor_width, floor_length), dtype=int)

for f in range(num_floors):
    for i in range(wall_thickness):
        floors[f,0+i,:] = 1; floors[f,-1-i,:] = 1
        floors[f,:,0+i] = 1; floors[f,:,-1-i] = 1

# 随机布置上楼梯和下楼梯（不同位置）
def place_stairs(num_up):
    up_list = []
    flag = np.zeros((floor_width, floor_length), dtype=int)
    for _ in range(num_up):
        while True:
            i,j = random.randint(wall_thickness,floor_width-wall_thickness-staircase_width), random.randint(wall_thickness,floor_length-wall_thickness-staircase_width)
            if sum(sum(flag[i:i+wall_thickness, j:j+wall_thickness])) == 0:
                for add_i in range(wall_thickness):
                    for add_j in range(wall_thickness):
                        up_list.append((i+add_i,j+add_j))
                flag[i:i+wall_thickness, j:j+wall_thickness] = 1
                break
    return up_list

downs = None
for f in range(num_floors):
    if f != num_floors - 1:
        ups = place_stairs(stairs_per_floor)
        for i,j in ups:
            floors[f,i,j] = 2  # 2代表上楼梯        
    if f != 0:
        for i,j in downs:
            floors[f,i,j] = 3  # 3代表下楼梯        
    downs = ups

# 随机放障碍物
for f in range(num_floors):
    for i in range(obstacle_attempt):
        if random.random() <= 0.5:
            h, w = obstacle_size
        else:
            w, h = obstacle_size
        rand_i = random.randint(wall_thickness+1,floor_width-wall_thickness-h-1)
        rand_j = random.randint(wall_thickness+1,floor_length-wall_thickness-w-1)            
        if sum(sum(floors[f,rand_i:rand_i+h,rand_j:rand_j+w])) == 0:
            floors[f,rand_i:rand_i+h,rand_j:rand_j+w] = 1

# 一楼随机放出口，只能在边界
exits = []
for _ in range(exits_on_ground):
    while True:
        side = random.choice(['top','bottom','left','right'])
        if side == 'top':
            pos = random.randint(wall_thickness,floor_length-wall_thickness-exit_width)
            if sum(floors[0, 0, pos:pos+exit_width]) == exit_width:
                floors[0, 0, pos:pos+exit_width] = 4   # 4代表出口
                floors[0, 1:wall_thickness, pos:pos+exit_width] = 0 # 0代表空地
                break
        elif side == 'bottom':
            pos = random.randint(wall_thickness,floor_length-wall_thickness-exit_width)
            if sum(floors[0, -1, pos:pos+exit_width]) == exit_width:
                floors[0, -1, pos:pos+exit_width] = 4   # 4代表出口
                floors[0, -wall_thickness:-1, pos:pos+exit_width] = 0 # 0代表空地
                break
        elif side == 'left':
            pos = random.randint(wall_thickness,floor_width-wall_thickness-exit_width)
            if sum(floors[0, pos:pos+exit_width, 0]) == exit_width:
                floors[0, pos:pos+exit_width, 0] = 4   # 4代表出口
                floors[0, pos:pos+exit_width, 1:wall_thickness] = 0 # 0代表空地 
                break
        else:
            pos = random.randint(wall_thickness,floor_width-wall_thickness-exit_width)
            if sum(floors[0, pos:pos+exit_width, -1]) == exit_width:
                floors[0, pos:pos+exit_width, -1] = 4   # 4代表出口
                floors[0, pos:pos+exit_width, -wall_thickness:-1] = 0 # 0代表空地  
                break

building = Building(num_floors=num_floors, 
                   floor_width=floor_width, 
                   floor_length=floor_length,
                   floor_height=floor_height,
                   cell_size=cell_size,
                   speed_floor_down_radio=speed_floor_down_radio,
                   floors=floors)

with open("data/Building_generate_0430.pkl", "wb") as haddle:
    pickle.dump(building, haddle)

# 画建筑物结构
fig, axes = plt.subplots(1, num_floors, figsize=(5*num_floors,5))
plt.tight_layout()

for f in range(num_floors):
    ax = axes[f]
    ax.set_title(f'Floor {f+1}')
    ax.set_xticks([]); ax.set_yticks([])
    np.where 
    # 障碍物或墙灰色方块
    pos_s = np.where(floors[f] == 1)       # 1代表障碍物或墙
    for i, j in zip(pos_s[0], pos_s[1]):
        i = int(i)
        j = int(j)
        ax.plot(i, j, 's', color='#808080', markersize=3, label='Wall or Obstable')
    # 上楼梯口蓝色标记
    pos_s = np.where(floors[f] == 2)       # 2代表上楼梯 
    for i, j in zip(pos_s[0], pos_s[1]):
        i = int(i)
        j = int(j)    
        ax.plot(i, j, 'b^', markersize=3, label='Upstairs')
    # 下楼梯口绿色标记        
    pos_s = np.where(floors[f] == 3)       # 3代表下楼梯 
    for i, j in zip(pos_s[0], pos_s[1]):
        i = int(i)
        j = int(j)  
        ax.plot(i, j, 'gv', markersize=3, label='Downstairs')      
    # 出口红色方块
    pos_s = np.where(floors[f] == 4)       # 4代表出口
    for i, j in zip(pos_s[0], pos_s[1]):
        i = int(i)
        j = int(j)      
        ax.plot(i, j, 'rs', markersize=3, label='Exit')

plt.savefig("data/Building_generate_0430.jpg")
